import time
import numpy as np
import onnxruntime as ort
import onnx

# 加载原始模型
onnx_model_path = "/home/lurker/文档/playground/python/pytorch/mnist_model.onnx"
session = ort.InferenceSession(onnx_model_path)

# 启用图优化
session_options = ort.SessionOptions()
session_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_EXTENDED
optimized_session = ort.InferenceSession(onnx_model_path, session_options)

# 准备输入数据
dummy_input = np.random.randn(1, 1, 28, 28).astype(np.float32)

# 推理未优化模型
start_time = time.time()
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name
session.run([output_name], {input_name: dummy_input})
end_time = time.time()
print(f"未优化模型推理时间：{end_time - start_time:.6f} 秒")

# 推理优化后的模型
start_time = time.time()
optimized_session.run([output_name], {input_name: dummy_input})
end_time = time.time()
print(f"优化模型推理时间：{end_time - start_time:.6f} 秒")
